"""
基于 Pinecone 的 KNN 手写数字识别 Web 应用
每次推理都使用 Pinecone 云服务进行
注意：第一次推理可能会较慢，因为需要初始化连接
"""

import gradio as gr
import numpy as np
from PIL import Image
import cv2
from pinecone import Pinecone
import os
import time

# 初始化 Pinecone
API_KEY = os.getenv("PINECONE_API_KEY", "pcsk_C1E8x_PBoUu1QoydYsUK4K5SWe5EDQvNskevRHiCpc7n8ozSny7Axtf99Ewm2F9LR9F7S")
INDEX_NAME = "mnist-digits-knn"
K_VALUE = 11  # KNN 中的 k 值

# 全局变量用于缓存连接
_pinecone_client = None
_pinecone_index = None
_first_inference = True


def get_pinecone_index():
    """获取 Pinecone 索引（懒加载）"""
    global _pinecone_client, _pinecone_index
    
    if _pinecone_index is None:
        print("初始化 Pinecone 连接...")
        _pinecone_client = Pinecone(api_key=API_KEY)
        _pinecone_index = _pinecone_client.Index(INDEX_NAME)
        print("Pinecone 连接初始化完成！")
    
    return _pinecone_index


def _to_numpy(image):
    """确保输入为 numpy 数组灰度图（uint8）。"""
    if image is None:
        return None
    if isinstance(image, dict):
        if 'image' in image:
            image = image['image']
        elif 'composite' in image:
            image = image['composite']
        else:
            image = image
    if isinstance(image, Image.Image):
        image = np.array(image)
    if not isinstance(image, np.ndarray):
        image = np.array(image)
    # If RGBA or RGB
    if image.ndim == 3:
        if image.shape[2] == 4:
            # use alpha channel as ink
            alpha = image[:, :, 3]
            gray = 255 - alpha
        else:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        gray = image
    return gray.astype(np.uint8)


def preprocess_for_digits(gray: np.ndarray) -> np.ndarray:
    """把任意手写图像转换为与 sklearn digits 数据集一致的特征向量 (1,64).

    流程：灰度->反转（若背景白）->阈值->裁剪->居中->模糊->缩放到8x8->归一化到0..16
    """
    # 自动反转（如果背景偏白）
    if np.mean(gray) > 127:
        gray = 255 - gray

    # 二值化（Otsu）以便更好地找到边界
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 找到非零像素位置
    coords = np.column_stack(np.where(thresh > 0))
    if coords.size == 0:
        return None

    y0, x0 = coords.min(axis=0)
    y1, x1 = coords.max(axis=0)

    # 添加小边距
    pad = 8
    y0 = max(0, y0 - pad)
    x0 = max(0, x0 - pad)
    y1 = min(gray.shape[0], y1 + pad)
    x1 = min(gray.shape[1], x1 + pad)

    roi = gray[y0:y1, x0:x1]

    # 放到正方形画布上并居中
    h, w = roi.shape
    size = max(h, w)
    square = np.zeros((size, size), dtype=np.uint8)
    y_off = (size - h) // 2
    x_off = (size - w) // 2
    square[y_off:y_off+h, x_off:x_off+w] = roi

    # 计算质心并将其移动到中心
    M = cv2.moments(square)
    if M['m00'] != 0:
        cx = M['m10'] / M['m00']
        cy = M['m01'] / M['m00']
        shift_x = (size / 2) - cx
        shift_y = (size / 2) - cy
        T = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
        square = cv2.warpAffine(square, T, (size, size), borderValue=0)

    # 平滑降噪，帮助 KNN 筛除细小噪点
    square = cv2.GaussianBlur(square, (3, 3), 0)

    # 缩放到 8x8
    small = cv2.resize(square, (8, 8), interpolation=cv2.INTER_AREA)

    # 归一化到 0..16（digits 数据集范围）
    small = small.astype(np.float32)
    # 对比度拉伸到 0..255 再映射到 0..16
    if small.max() > small.min():
        small = (small - small.min()) / (small.max() - small.min()) * 255.0
    small = small / 255.0 * 16.0

    return small.flatten()


def predict_digit_pinecone(image) -> str:
    """使用 Pinecone 进行预测，返回单个预测结果（字符串）"""
    global _first_inference
    
    try:
        # 显示第一次推理提示
        if _first_inference:
            print("=" * 60)
            print("注意：第一次推理可能需要几秒钟来初始化 Pinecone 连接")
            print("=" * 60)
            _first_inference = False
        
        start_time = time.time()
        
        gray = _to_numpy(image)
        if gray is None:
            return ""

        # 预处理图像
        features = preprocess_for_digits(gray)

        # 如果严格预处理失败，尝试回退方案
        if features is None:
            arr = gray.astype(np.uint8)
            if np.mean(arr) > 127:
                arr = 255 - arr
            small = cv2.resize(arr, (8, 8), interpolation=cv2.INTER_AREA).astype(np.float32)
            if small.max() > small.min():
                small = (small - small.min()) / (small.max() - small.min()) * 255.0
            small = small / 255.0 * 16.0
            features = small.flatten()

        # 获取 Pinecone 索引
        index = get_pinecone_index()
        
        # 查询 Pinecone
        query_results = index.query(
            vector=features.tolist(),
            top_k=K_VALUE,
            include_metadata=True
        )
        
        # 收集邻居标签
        neighbor_labels = []
        for match in query_results['matches']:
            if 'metadata' in match and 'label' in match['metadata']:
                neighbor_labels.append(match['metadata']['label'])
        
        # 多数投票
        if not neighbor_labels:
            return ""
        
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        
        elapsed_time = time.time() - start_time
        print(f"推理完成，耗时: {elapsed_time:.3f} 秒")
        
        return str(int(predicted_label))
        
    except Exception as e:
        import traceback
        print("=" * 60)
        print("预测出错:")
        traceback.print_exc()
        print("=" * 60)
        print("\n请确保：")
        print("1. 已设置正确的 PINECONE_API_KEY 环境变量")
        print("2. 已运行 pinecone_train.py 创建索引")
        print("3. 索引名称为:", INDEX_NAME)
        return "错误"


def clear_all():
    """清空画布和结果"""
    return None, ""


# 创建 Gradio 界面
with gr.Blocks(title="Pinecone KNN 手写数字识别") as demo:
    gr.Markdown("""
    # Pinecone KNN 手写数字识别
    
    在左侧画布写数字，右侧显示预测结果。
    
    **注意：** 
    - 本应用使用 Pinecone 云服务进行推理
    - 第一次推理可能需要几秒钟来初始化连接
    - 后续推理会快很多
    - 请确保已运行 `pinecone_train.py` 创建索引
    """)

    with gr.Row():
        with gr.Column(scale=2):
            # 兼容不同版本的 Gradio
            try:
                sketch = gr.Sketchpad(
                    image_mode="RGB",
                    canvas_size=(280, 280),
                    label="手写区"
                )
            except Exception:
                sketch = gr.Image(
                    source="canvas",
                    tool="sketch",
                    type="pil",
                    label="手写区"
                )
            
            with gr.Row():
                predict_btn = gr.Button("识别", variant="primary")
                clear_btn = gr.Button("清除")
        
        with gr.Column(scale=1):
            result = gr.Textbox(
                label="预测结果",
                lines=1,
                interactive=False
            )
            gr.Markdown("""
            ### 使用说明
            1. 在画布上手写一个数字（0-9）
            2. 点击"识别"按钮
            3. 等待几秒（第一次会较慢）
            4. 查看预测结果
            5. 点击"清除"可以重新开始
            """)

    predict_btn.click(
        fn=predict_digit_pinecone,
        inputs=[sketch],
        outputs=[result]
    )
    clear_btn.click(
        fn=clear_all,
        inputs=None,
        outputs=[sketch, result]
    )


if __name__ == '__main__':
    print("\n" + "=" * 60)
    print("启动 Pinecone KNN 手写数字识别 Web 应用")
    print("=" * 60)
    print(f"使用索引: {INDEX_NAME}")
    print(f"K 值: {K_VALUE}")
    print("=" * 60 + "\n")
    
    demo.launch(share=False)
